#include "win_util.h"
#include <windows.h>
#include "log.h"
#include "string_util.h"

namespace {
constexpr DWORD MAX_NAME_LEN = 4096;
constexpr int MIN_PATH_LEN = 2;

const std::vector<std::wstring> UNINSTALL_REG_ROOT_LIST = {
    L"SOFTWARE\\WOW6432Node\\Microsoft\\Windows\\CurrentVersion\\Uninstall",
    L"SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Uninstall"};
const std::vector<HKEY> HKEY_LIST = {HKEY_LOCAL_MACHINE, HKEY_CURRENT_USER};
const std::wstring VAR_TRANSLATION_NAME = L"\\VarFileInfo\\Translation";
const std::string VAR_PRODUCT_FORMATE = "\\StringFileInfo\\{:04x}{:04x}\\ProductName";
const std::vector<std::wstring> REG_NAME_LIST = {L"InstallLocation", L"UninstallString", L"DisplayIcon"};

bool ReadRegVal(HKEY key, const std::wstring &name, std::wstring &val)
{
    DWORD size = MAX_NAME_LEN;

    auto data = std::make_unique<BYTE[]>(size);
    if (data == nullptr) {
        LOG_ERR("create buff failed.");
        return false;
    }
    SecureZeroMemory(data.get(), size);
    DWORD readSize = size;
    DWORD dwType = REG_SZ;
    DWORD result = 0;
    if ((result = RegQueryValueExW(key, name.c_str(), nullptr, &dwType, data.get(), &readSize)) != ERROR_SUCCESS) {
        LOG_DEBUG("read reg val failed. err: {}", result);
        return false;
    }
    val = std::wstring(reinterpret_cast<wchar_t *>(data.get()));
    return true;
}

std::vector<BYTE> GetFileVersionInfoData(const std::wstring &filePath)
{
    DWORD handle;
    DWORD versionInfoSize = GetFileVersionInfoSizeW(filePath.c_str(), &handle);
    if (versionInfoSize == 0) {
        LOG_WARN("Error: Unable to retrieve version information size.");
        return {};
    }
    std::vector<BYTE> versionData(versionInfoSize);
    if (!GetFileVersionInfoW(filePath.c_str(), handle, versionInfoSize, versionData.data())) {
        LOG_WARN("Error: Unable to retrieve version information.");
        return {};
    }
    return versionData;
}

bool ReadFileProdName(const std::wstring &filePath, std::wstring &productStr)
{
    std::vector<BYTE> versionData = GetFileVersionInfoData(filePath);
    if (versionData.empty()) {
        LOG_ERR("get file version info failed, err: {}", GetLastError());
        return false;
    }
    struct LANGANDCODEPAGE {
        WORD wLanguage;
        WORD wCodePage;
    } *lpTranslate;
    UINT translateSize;
    auto ret = VerQueryValueW(
        versionData.data(), VAR_TRANSLATION_NAME.c_str(), reinterpret_cast<LPVOID *>(&lpTranslate), &translateSize);
    if (!ret) {
        LOG_WARN("Error: Unable to read translation info data.");
        return false;
    }
    for (int i = 0; i < translateSize / sizeof(LANGANDCODEPAGE); i++) {
        LPVOID queryResult;
        UINT dwBytes;
        auto varProductQuery =
            Utf8ToUnicode(fmt::format(VAR_PRODUCT_FORMATE, lpTranslate[i].wLanguage, lpTranslate[i].wCodePage));
        if (!VerQueryValueW(versionData.data(), varProductQuery.c_str(), &queryResult, &dwBytes)) {
            LOG_WARN("Error: Unable to read product info data.");
            continue;
        }
        std::wstring productName = static_cast<PWCHAR>(queryResult);
        if (!productName.empty()) {
            productStr = productName;
            return true;
        }
    }
    return false;
}

std::vector<std::wstring> GetInstallPath(std::wstring &installPath, std::wstring &uninstallString)
{
    std::vector<std::wstring> ret;
    if (!installPath.empty()) {
        ret.push_back(installPath);
    }
    if (uninstallString.empty()) {
        return ret;
    }
    if (uninstallString.find(L"MsiExec.exe") != std::wstring::npos) {
        return ret;
    }
    if (uninstallString.find('\"') == std::wstring::npos) {
        uninstallString = uninstallString.substr(0, uninstallString.rfind('\\'));
        if (ret.empty()) {
            ret.push_back(uninstallString);
            return ret;
        }
    }
    if (uninstallString.size() > MIN_PATH_LEN) {
        size_t start = uninstallString.find('\"');
        size_t end = uninstallString.rfind('\\');
        if (end > start) {
            uninstallString = uninstallString.substr(start + 1, end - start);
        }
    }
    if (!installPath.empty() && uninstallString.find(installPath) == 0) {
        return ret;
    }
    ret.push_back(uninstallString);
    return ret;
}

bool OpenRegKey(HKEY field, const std::wstring &root, PHKEY key, bool isReadOnly = true)
{
    auto flag = KEY_READ;
    if (!isReadOnly) {
        flag = KEY_ALL_ACCESS;
    }
    int ret = RegOpenKeyExW(field, root.c_str(), 0, flag, key);
    if (ret != ERROR_SUCCESS) {
        LOG_ERR("open reg key failed.");
        return false;
    }
    return true;
}

bool IsIncludePath(std::vector<std::wstring> &pathList, const std::wstring &path)
{
    for (const auto &item : pathList) {
        if (path.find(item) == 0) {
            return true;
        }
    }
    return false;
}

bool ParsePath(const std::wstring &source, std::wstring &path)
{
    if (source.empty()) {
        return false;
    }
    if (source.find(L"MsiExec.exe") != std::wstring::npos) {
        return false;
    }
    path = source;
    struct stat s {};
    if (stat(UnicodeToUtf8(source).c_str(), &s) == 0 && (s.st_mode & S_IFDIR)) {
        if (path.rfind(L'\\') == path.length() - 1) {
            path = path.substr(0, path.length() - 1);
        }
        return true;
    }
    size_t start = path.find('\"');
    if (start == 0) {
        path = path.substr(1);
    }
    size_t end = path.rfind('\\');
    if (end != std::wstring::npos) {
        path = path.substr(0, end);
    }
    return !path.empty();
}

std::vector<std::wstring> GetInstalledPathList(HKEY filed, const std::wstring &root)
{
    std::vector<std::wstring> pathList;
    HKEY key;
    if (!OpenRegKey(filed, root, &key)) {
        LOG_ERR("open reg failed, err: {}", GetLastError());
        return pathList;
    }
    DWORD index = 0;
    DWORD size = MAX_NAME_LEN;
    auto buf = std::make_unique<wchar_t[]>(size);
    if (buf == nullptr) {
        LOG_ERR("create buf failed. err: {}", GetLastError());
        return pathList;
    }
    SecureZeroMemory(buf.get(), sizeof(wchar_t) * size);
    DWORD result = 0;
    while (
        (result = RegEnumKeyExW(key, index++, buf.get(), &size, nullptr, nullptr, nullptr, nullptr)) == ERROR_SUCCESS) {
        HKEY subItemKey;
        std::wstring subItemName(buf.get());
        if (!OpenRegKey(key, subItemName, &subItemKey)) {
            LOG_ERR("open key [{}] failed, err: {}", WstringToString(subItemName), GetLastError());
            size = MAX_NAME_LEN;
            SecureZeroMemory(buf.get(), sizeof(wchar_t) * size);
            continue;
        }
        for (auto &item : REG_NAME_LIST) {
            std::wstring source;
            std::wstring path;
            if (!ReadRegVal(subItemKey, item, source)) {
                continue;
            }
            if (ParsePath(source, path) && !IsIncludePath(pathList, path)) {
                pathList.push_back(path);
            }
        }

        SecureZeroMemory(buf.get(), sizeof(wchar_t) * size);
        RegCloseKey(subItemKey);
        size = MAX_NAME_LEN;
    }
    if (result != ERROR_NO_MORE_ITEMS) {
        LOG_ERR("enum failed. err: {}", result);
    }
    RegCloseKey(key);
    return pathList;
}
}  // namespace

std::wstring GetProcessName(DWORD pid)
{
    std::wstring ret;
    if (HANDLE handle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid)) {
        DWORD buffSize = MAX_NAME_LEN;
        auto buf = std::make_unique<wchar_t[]>(MAX_NAME_LEN);
        if (buf == nullptr) {
            LOG_ERR("create buf failed. err: {}", GetLastError());
            CloseHandle(handle);
            return ret;
        }
        SecureZeroMemory(buf.get(), sizeof(wchar_t) * buffSize);
        if (QueryFullProcessImageNameW(handle, 0, buf.get(), &buffSize)) {
            ret = buf.get();
        } else {
            LOG_ERR("get process failed. err: {}", GetLastError());
        }
        CloseHandle(handle);
    } else {
        LOG_ERR("Error OpenProcess : {}", GetLastError());
    }
    return ret;
}

std::vector<std::wstring> GetInstalledPathList()
{
    std::vector<std::wstring> pathList;
    for (auto key : HKEY_LIST) {
        for (auto reg : UNINSTALL_REG_ROOT_LIST) {
            for (auto &path : GetInstalledPathList(key, reg.c_str())) {
                pathList.push_back(path);
            }
        }
    }
    return pathList;
}

std::shared_ptr<ProductInfo> GetProcessFileProductInfo(const std::wstring &filePath)
{
    std::unique_ptr<ProductInfo> info = nullptr;
    std::wstring productName;
    if (ReadFileProdName(filePath, productName)) {
        info = std::make_unique<ProductInfo>();
        if (info != nullptr) {
            info->productName = productName;
        }
    }
    return info;
}

std::wstring GetFileName(const std::wstring &path)
{
    size_t pos = path.rfind(L'\\');
    if (pos == std::wstring::npos) {
        return path;
    }
    return path.substr(pos + 1, path.size());
}
